import os
import uuid
from itertools import groupby
from typing import List, Optional, Dict

import numpy as np
import numba as nb

from common.logger import logger
from common.config import DEFAULT_CACHE_DIR
from common.typing_utils import numeric
from common.io_utils.dirs import mkdir_p
from common.io_utils.temps import temp_dir

from common.geometry.geom import Geom
from common.geometry.bbox import BoundingBox
from common.geometry.tools import project_geom
from common.geoimage.gdal_sys import mosaic_ds  # , warp_ds
from common.geoimage.raster_dataset import RasterDataset

from catalog.models.base_stac_item import BaseSTACItem
from catalog.raster.utils import (
    determine_dest_srs,
    determine_dest_res,
    determine_dest_extent,
    is_s2_l2a_collection,
    update_s2_l2a_boa,
    get_s2_l2a_add_boa_offset,
    nb_customized_median,
)


# NOTE: sentinel-2 L2A SCL data dictionary
# Sen2-Cor-L2A-Input-Output-Data-Definition-Document, Page 24-25, Table 8
# 0: nodata; 3: cloud shadow, 8: cloud medium prob, 9: cloud high prob, 11: snow or ice
# 10: thin cirrus, 7: unclassified, 6: water, 5: not veg, 4: veg, 2: dark features, shadows
VALS_TO_MASK = [0, 3, 8, 9, 11]

TMP_DIR = os.path.join(DEFAULT_CACHE_DIR, "catalog", "raster")
mkdir_p(TMP_DIR)


@nb.jit(nb.uint16[:, :](nb.uint16[:, :], nb.uint16[:, :, :]), nopython=True, cache=True)
def nb_customized_fillnodata(nda_res, nda_values):
    nsrcs, nrows, ncols = nda_values.shape
    nda_tmp = np.zeros(nsrcs, dtype=np.uint16)
    for i in range(nrows):
        for j in range(ncols):
            if nda_res[i, j] == 0:
                available_values = 0
                for k in range(nsrcs):
                    if nda_values[k, i, j] != 0:
                        nda_tmp[available_values] = nda_values[k, i, j]
                        available_values += 1
                if available_values > 0:
                    nda_res[i, j] = np.median(nda_tmp[0:available_values])
    return nda_res


def composite_s2decloud_v1_per_mgrs_tile(
    items: List[BaseSTACItem], asset_names: Dict[str, str], aoi: Geom
):
    """composite a list of items in the same MGRS tile, using s2decloud-v1 method"""
    if len(items) == 1:
        item = items[0]
        vrt_scene = item.choose_assets(asset_names)
        rds = RasterDataset.from_vrt_scene(vrt_scene, aoi=aoi)
        # NOTE Post-process Sentinel-2 L2A due to ESA updates of -1000 offset
        asset_offsets = get_s2_l2a_add_boa_offset(item, asset_names)
        if asset_offsets is not None:
            rds.data = update_s2_l2a_boa(rds.data, asset_offsets)
        resolution = abs(rds.meta.x_res)

        vrt_scene_scl = item.choose_assets(["SCL"]).warp(
            dest_x_res=resolution, dest_y_res=resolution
        )
        rds_scl = RasterDataset.from_vrt_scene(vrt_scene_scl, aoi=aoi)
        scl_mask = np.isin(np.squeeze(rds_scl.data), VALS_TO_MASK)

        meta = rds.meta
        nda_data = []
        # TODO get a numba based faster implementation
        for band_ind in range(meta.n_bands):
            nda_data.append(np.where(scl_mask, 0, rds.data[band_ind, :, :]))
        nda_data = np.stack(nda_data, axis=0)
        return RasterDataset.from_ndarray(nda_data, meta)

    loaded_rds = []
    for item in items:
        vrt_scene = item.choose_assets(asset_names)
        rds = RasterDataset.from_vrt_scene(vrt_scene, aoi=aoi)
        # NOTE Post-process Sentinel-2 L2A due to ESA updates of -1000 offset
        asset_offsets = get_s2_l2a_add_boa_offset(item, asset_names)
        if asset_offsets is not None:
            rds.data = update_s2_l2a_boa(rds.data, asset_offsets)
        loaded_rds.append(rds)
    resolution = abs(rds.meta.x_res)

    loaded_rds_scl = []
    for item in items:
        vrt_scene_scl = item.choose_assets(["SCL"]).warp(
            dest_x_res=resolution, dest_y_res=resolution
        )
        rds_scl = RasterDataset.from_vrt_scene(vrt_scene_scl, aoi=aoi)
        loaded_rds_scl.append(rds_scl)

    stacked_scl_mask = []
    for rds_scl_tmp in loaded_rds_scl:
        scl_mask = np.isin(np.squeeze(rds_scl_tmp.data), VALS_TO_MASK)
        stacked_scl_mask.append(scl_mask)
    stacked_scl_mask = np.stack(stacked_scl_mask).astype(np.uint8)

    nda_res = []
    for ind_band, _ in enumerate(asset_names):
        nda_data_tmp = np.stack(
            [np.squeeze(rds_tmp.data[ind_band, :, :]) for rds_tmp in loaded_rds]
        )
        nda_res.append(nb_customized_median(nda_data_tmp, stacked_scl_mask))
    nda_res = np.stack(nda_res)

    return RasterDataset.from_ndarray(nda_res, loaded_rds[-1].meta)


def composite_s2decloud_v1(
    items: List[BaseSTACItem],
    asset_names: List[str],
    extent: Optional[BoundingBox] = None,
    aoi: Optional[Geom] = None,
    dest_srs: Optional[str] = None,
    dest_res: Optional[numeric] = None,
    resample_alg: str = "nearest",
):
    """Composite a given list of `BaseSTACItems` using a median composite method,
    removing cloud, cloud shadow, and snow-ice.

    NOTE: This function is designed to be used within `catalog.raster.raster` only.
    It relies on the caller to provide several arguments correctly.
    """
    if not is_s2_l2a_collection(items):
        raise ValueError(
            (
                "The input items are not from the COPERNICUS/SENTINEL2/L2A product"
                " and not supported by this composite strategy."
            )
        )

    if dest_srs is None:
        dest_srs = determine_dest_srs(aoi, items)

    if dest_res is None:
        dest_res = determine_dest_res(dest_srs, items)

    dest_extent = determine_dest_extent(extent, aoi, dest_srs, dest_res)

    # NOTE derive a `work_aoi`, a special part for the `s2decloud-v1` composite strategy
    # use that `work_aoi` to make per-tile process more efficient by only working on needed subset
    if dest_extent is not None:
        # if an extent is given, convert it to a geometry under WGS84 as the AOI
        bbox_geom = Geom.from_bbox(dest_extent)
        if dest_srs.upper() != "EPSG:4326":
            work_aoi = project_geom(bbox_geom, src_epsg=dest_srs, dest_epsg="EPSG:4326")
        else:
            work_aoi = bbox_geom
    else:
        work_aoi = None

    # group by mgrs tile
    items = sorted(items, key=lambda item: item.mgrs_tile)
    tile_groups = groupby(items, key=lambda item: item.mgrs_tile)
    n_tiles = len(set(item.mgrs_tile for item in items))

    if n_tiles == 1:
        for mgrs_tile, sub_items in tile_groups:
            sub_items = list(sub_items)
        logger.info(f"working on {mgrs_tile} with {len(sub_items)} images")
        rds = composite_s2decloud_v1_per_mgrs_tile(
            sub_items,
            asset_names,
            work_aoi,
        )
        rds = rds.warp(dest_srs=dest_srs, dest_res=dest_res, dest_extent=dest_extent)
        return rds

    per_tile_files = []
    temp_dir_name = str(uuid.uuid4())
    with temp_dir(
        dir_name=temp_dir_name, base_dir=TMP_DIR, remove=True
    ) as temp_dir_path:
        for mgrs_tile, sub_items in tile_groups:
            sub_items = list(sub_items)
            logger.info(f"working on {mgrs_tile} with {len(sub_items)} images")
            rds = composite_s2decloud_v1_per_mgrs_tile(
                sub_items,
                asset_names,
                work_aoi,
            )
            if dest_srs != rds.meta.prj_epsg:
                rds = rds.warp(dest_srs=dest_srs)
            fpath_tmp = os.path.join(temp_dir_path, f"{mgrs_tile}.tif")
            rds.to_geotiff(fpath_tmp)
            per_tile_files.append(fpath_tmp)

        with mosaic_ds(
            per_tile_files,
            resample_alg=resample_alg,
            dest_res=dest_res,
            dest_extent=dest_extent,
            src_nodata=0,
        ) as gdal_rst:
            rds = RasterDataset.from_gdal_raster(gdal_rst, extent=dest_extent)

    return rds
